SIMD 示例

所谓 SIMD 就是一次指令计算多个数据，例如 AVX256 一次计算 256 位数据。

  • int 是 32 位，所以 AVX256 一次计算 8 个
  • double 是 64 位，所以一次计算 4 个

以计算 double 加法为例:

1__m256d m256x; // 定义标识 AVX 寄存器的变量
2__m256d m256y;
3
4m256x = _mm256_set_pd(x[i+3], x[i+2], x[i+1], x[i]); // 向 AVX 寄存器中写入数据，大端序
5m256y = _mm256_set_pd(y[i+3], y[i+2], y[i+1], y[i]);
6m256x = _mm256_add_pd(m256x, m256y);                 // 一次操作 4 对 double 数据的加法
7
8double o[4] __attribute__((aligned(32)));
9_mm256_store_pd(o, m256x);                           // 读出结果，必须 32 位对齐

使用 gcc 编译时，需要附带 -mavx2 选项。

在 Linux 上，执行 cat /proc/cpuinfo | grep flags 命令可以查看 CPU 支持哪些特性:

1$ cat /proc/cpuinfo | grep flags
2flags		: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat 
3pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp 
4lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni 
5pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx 
6f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 
73dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext 
8perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp 
9vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap 
10clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc 
11cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd arat npt 
12lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists 
13pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip pku ospke 
14vaes vpclmulqdq rdpid overflow_recov succor smca fsrm
名称说明
MMX64位整型
SSE128位单精度浮点
SSE2扩展128位整型、双精度浮点，CPU快取的控制指令
SSE3扩展类型转换，超线程支持
SSE4扩展 CRC32 等指令
AVX256位浮点
AVX2扩展256位整型，三操作数指令（3-Operand Instructions）
AVX512512位运算
  • MMX 是 MultiMedia eXtensions 的缩写
  • SSE 是 Streaming SIMD Extensions 的缩写
  • AVX 是 Advanced Vector Extensions 的缩写

注意不要混用不同的指令（例如同时使用 AVX 和 SSE），否则将会引发 transition penalty，造成性能损失。

在 gcc 中需要引用 x86intrin.h 头文件，在 MSVC 中需要引用 intrin.h 头文件。

具体的:

1#include <mmintrin.h>   // mmx, 4个64位寄存器
2#include <xmmintrin.h>  // sse, 8个128位寄存器
3#include <emmintrin.h>  // sse2, 8个128位寄存器
4#include <pmmintrin.h>  // sse3, 8个128位寄存器
5#include <smmintrin.h>  // sse4.1, 8个128位寄存器
6#include <nmmintrin.h>  // sse4.2, 8个128位寄存器
7#include <immintrin.h>  // avx, 16个256位寄存器

示例程序 demo1.c 中，可以看到使用 avx256 和直接计算，耗时差不多，甚至直接计算反而更快。这是因为运算过于简单，数据拷贝占用的时间更高。 因此使用 SIMD 的时候，需要把算法的步骤组合起来，在一次 IO 中进行尽可能多的运算。

另外还发现，在 AMD Ryzen 9 5900HX 上使用 gcc 编译，优化级别设为 -O0 时，AVX256 会比直接计算慢很多。 不清楚是 CPU 的原因还是编译器的原因。

示例1

1#include <stdio.h>
2#include <string.h>
3#include <x86intrin.h>
4#include <time.h>
5
6void avx256_vector_add(const double* x, const double* y, size_t n, double* out);
7void avx256_vector_sub(const double* x, const double* y, size_t n, double* out);
8void avx256_vector_mul(const double* x, const double* y, size_t n, double* out);
9void avx256_vector_div(const double* x, const double* y, size_t n, double* out);
10
11void vector_add(const double* x, const double* y, size_t n, double* out);
12void vector_sub(const double* x, const double* y, size_t n, double* out);
13void vector_mul(const double* x, const double* y, size_t n, double* out);
14void vector_div(const double* x, const double* y, size_t n, double* out);
15
16#define TEST(func, ...) do{                                             \
17                            clock_t start = clock();                    \
18                            func(__VA_ARGS__);                          \
19                            clock_t end = clock();                      \
20                            printf(#func " cost: %ld\n", end - start);  \
21                        }while(0)
22
23int main(void)
24{
25    const size_t N = 2048;
26    double x[N];
27    double y[N];
28    double out[N];
29
30    TEST(avx256_vector_add, x, y, N, out);
31    TEST(avx256_vector_sub, x, y, N, out);
32    TEST(avx256_vector_mul, x, y, N, out);
33    TEST(avx256_vector_div, x, y, N, out);
34
35    TEST(vector_add, x, y, N, out);
36    TEST(vector_sub, x, y, N, out);
37    TEST(vector_mul, x, y, N, out);
38    TEST(vector_div, x, y, N, out);
39}
40
41void avx256_vector_add(const double* x, const double* y, size_t n, double* out)
42{
43    __m256d m256x;
44    __m256d m256y;
45    double o[4] __attribute__((aligned(32)));
46
47    size_t i = 0;
48    for (; i + 3 < n; i+=4)
49    {
50        m256x = _mm256_set_pd(x[i+3], x[i+2], x[i+1], x[i]);
51        m256y = _mm256_set_pd(y[i+3], y[i+2], y[i+1], y[i]);
52        m256x = _mm256_add_pd(m256x, m256y);
53        _mm256_store_pd(o, m256x);
54
55        memcpy(out, o, sizeof(double) * 4);
56    }
57
58    // 剩下少量没对齐的数据没必要使用 SIMD
59    for(;i < n; i++)
60    {
61        out[i] = x[i] + y[i];
62    }
63}
64
65void avx256_vector_sub(const double* x, const double* y, size_t n, double* out)
66{
67    __m256d m256x;
68    __m256d m256y;
69    double o[4] __attribute__((aligned(32)));
70
71    size_t i = 0;
72    for (; i + 3 < n; i+=4)
73    {
74        m256x = _mm256_set_pd(x[i+3], x[i+2], x[i+1], x[i]);
75        m256y = _mm256_set_pd(y[i+3], y[i+2], y[i+1], y[i]);
76        m256x = _mm256_sub_pd(m256x, m256y);
77        _mm256_store_pd(o, m256x);
78
79        memcpy(out, o, sizeof(double) * 4);
80    }
81
82    // 剩下少量没对齐的数据没必要使用 SIMD
83    for(;i < n; i++)
84    {
85        out[i] = x[i] - y[i];
86    }
87}
88
89void avx256_vector_mul(const double* x, const double* y, size_t n, double* out)
90{
91    __m256d m256x;
92    __m256d m256y;
93    double o[4] __attribute__((aligned(32)));
94
95    size_t i = 0;
96    for (; i + 3 < n; i+=4)
97    {
98        m256x = _mm256_set_pd(x[i+3], x[i+2], x[i+1], x[i]);
99        m256y = _mm256_set_pd(y[i+3], y[i+2], y[i+1], y[i]);
100        m256x = _mm256_mul_pd(m256x, m256y);
101        _mm256_store_pd(o, m256x);
102
103        memcpy(out, o, sizeof(double) * 4);
104    }
105
106    // 剩下少量没对齐的数据没必要使用 SIMD
107    for(;i < n; i++)
108    {
109        out[i] = x[i] * y[i];
110    }
111}
112
113void avx256_vector_div(const double* x, const double* y, size_t n, double* out)
114{
115    __m256d m256x;
116    __m256d m256y;
117    double o[4] __attribute__((aligned(32)));
118
119    size_t i = 0;
120    for (; i + 3 < n; i+=4)
121    {
122        m256x = _mm256_set_pd(x[i+3], x[i+2], x[i+1], x[i]);
123        m256y = _mm256_set_pd(y[i+3], y[i+2], y[i+1], y[i]);
124        m256x = _mm256_div_pd(m256x, m256y);
125        _mm256_store_pd(o, m256x);
126
127        memcpy(out, o, sizeof(double) * 4);
128    }
129
130    // 剩下少量没对齐的数据没必要使用 SIMD
131    for(;i < n; i++)
132    {
133        out[i] = x[i] / y[i];
134    }
135}
136
137void vector_add(const double* x, const double* y, size_t n, double* out)
138{
139    for (size_t i = 0; i < n; i++)
140    {
141        out[i] = x[i] + y[i];
142    }
143}
144void vector_sub(const double* x, const double* y, size_t n, double* out)
145{
146    for (size_t i = 0; i < n; i++)
147    {
148        out[i] = x[i] - y[i];
149    }
150}
151
152void vector_mul(const double* x, const double* y, size_t n, double* out)
153{
154    for (size_t i = 0; i < n; i++)
155    {
156        out[i] = x[i] * y[i];
157    }
158}
159
160void vector_div(const double* x, const double* y, size_t n, double* out)
161{
162    for (size_t i = 0; i < n; i++)
163    {
164        out[i] = x[i] / y[i];
165    }
166}

示例2

1#include <stdio.h>
2#include <string.h>
3#include <math.h>
4#include <x86intrin.h>
5#include <time.h>
6
7double avx256_pi(size_t limit);
8double pi(size_t limit);
9
10#define TEST(func, ...) do{                                             \
11                            clock_t start = clock();                    \
12                            func(__VA_ARGS__);                          \
13                            clock_t end = clock();                      \
14                            printf(#func " cost: %ld\n", end - start);  \
15                        }while(0)
16
17int main(void)
18{
19    #define N 0xfffff
20
21    TEST(pi, N);
22    printf("%f\n", pi(N));
23
24    TEST(avx256_pi, N);
25    printf("%f\n", avx256_pi(N));
26
27
28}
29
30double avx256_pi(size_t limit)
31{
32    // ∑(1 / n^2) = pi^2 / 6
33
34    __m256d sum = _mm256_set1_pd(0.0); // 四个 double 都设为同一个值 0.0
35    __m256d one = _mm256_set1_pd(1.0);
36    __m256d temp;
37
38    for (size_t n = 1; n + 3 < limit; n+=4)
39    {
40        temp = _mm256_set_pd(n+3.0, n+2.0, n+1.0, n);
41        temp = _mm256_mul_pd(temp, temp);
42        temp = _mm256_div_pd(one, temp);
43        sum = _mm256_add_pd(sum, temp);
44    }
45
46
47    double o[4] __attribute__((aligned(32)));
48    _mm256_store_pd(o, sum);
49
50    return sqrt(6*(o[0] + o[1] + o[2] + o[3]));
51}
52
53double pi(size_t limit)
54{
55    // // ∑(1 / n^2) = pi^2 / 6
56
57    double sum = 0.0;
58    for (size_t n = 1; n < limit; n++)
59    {
60        sum += 1.0 / n / n;
61    }
62
63    return sqrt(6 * sum);
64}